#! /usr/bin/env python
# 
#  run_mpi.py
#  Run a job on green easily from the command line from any node.
#  Based on csh script by Greg Poole (Swinburne University of Technology). 
#  
#  Created by Simon Mutch on 2010-09-23.
#  Email: smutch@astro.swin.edu.au
#  Copyright 2010 Swinburne University of Technology. All rights reserved.
# 
#  N.B. The author takes ZERO responsibility for crashing any supercomputers or any
#  other undesireable effects as a result of using this script.
# 

import sys
import getopt
from tempfile import NamedTemporaryFile
from os import getcwd
import subprocess as sp

# USER PARAMETERS - Set these to your desired values ========
user_params = {
                "set_email":True,
                "email_address":"smutch@astro.swin.edu.au",
                "mpi_dir":"/usr/local/gnu/x86_64/openmpi-1.4",
                "mpi_opt":"-mca btl tcp,self",
                "default_nodes":1,
                "default_ppn":1,
                "default_pmem":"1gb",
                "default_walltime":"00:59:59"                
}
# ===========================================================

def useage():
    
    """ Print the useage message. """
    
    print main.__doc__


def write_script(script, job_name, pwd, res, job, use_mpiexec):
    
    """ Actually write the submission script given the passed arguments. 
    
        This function should be modified to suit your needs.
    """

    email_address = "smutch@astro.swin.edu.au"

    # This is all pretty standard stuff
    script.write("""#!/bin/csh
#PBS -z
#PBS -S /bin/csh
#PBS -V
""")

    if user_params["set_email"]:
        script.write("#PBS -m ae\n#PBS -M "+user_params["email_address"]+"\n")
    
    # These lines are generated from the input provided on the command line
    script.write("#PBS -N "+job_name+"\n")
    script.write("#PBS -d "+pwd+"\n")
    script.write("#PBS -l "+res+"\n\n")

    # These are generated from the user_params dict
    if use_mpiexec == True:
        script.write("set PATH = "+user_params["mpi_dir"]+"/bin:$PATH\n")
        script.write("set LD_LIBRARY_PATH = "+user_params["mpi_dir"]+"/lib:$LD_LIBRARY_PATH\n")
    
    # These lines should not need to be changed
    script.write("""
cd $PBS_O_WORKDIR
                                              
echo '-------------------------------' 
echo "DATE     is "`date`     
echo "PWD      is "$PWD            
echo "NNODES   is "`uniq $PBS_NODEFILE | wc -l`         
echo "NPROCS   is "`cat  $PBS_NODEFILE | wc -l`         
echo "Using the following nodes:"      
cat $PBS_NODEFILE | uniq               
echo '-------------------------------' 

""")

    if use_mpiexec == True:
        # This assumes we are running an mpi job
        script.write(user_params["mpi_dir"]+"/bin/mpiexec "+user_params["mpi_opt"]+" "+job+" \n\n")
    else:
        # We don't want to use mpiexec in this case
        script.write(job+" \n\n")

    # The all important line - unless you want an empty submission script!
    # We do this because we want to do the rather dodgy process doing something with the file
    # (i.e. submitting it to the queue) before we close it.
    script.flush()



def main(argv):
    
    """
    Generate a script for a program and submit it to the queue on the supercomputer.

    Useage:
    run_mpi [options] [command to run]

    Options:
    -r, (--resources=) -- set the resources part of the job script.
                          e.g. nodes=1:ppn=2:s203,walltime=23:59:59
    -g, (--gen_script=) -- generate the script using the provided file name 
                          only - do not submit it to the queue.
    -n, (--job_name=)   -- the name to be given to the job on the queue.
    -h, (--help)       -- print this help message.
    --no_mpiexec       -- don't prepend the default mpiexec to the job

    Notes:
    If you do not specify any of nodes, ppn, walltime or pmem then defaults
    will be assumed.  Modify this script directly to change what these defaults are.

    This script, in it's present form, may be run from any node on green.  For best
    results have your ssh agent up and running.

    Example:
    run_mpi -r nodes=1:ppn=2:s203 --job_name=test_submission program arg1 arg2 arg3
    """

    from optparse import OptionParser

    version = "0.2"

    description = "Generate a script for a program and submit it to the queue on the supercomputer."

    usage = "usage: %prog [options] executable [args]"

    epilog = """Notes: If you do not specify any of nodes, ppn, walltime or pmem then defaults
will be assumed.  Modify this script directly to change what these defaults are.
This script, in it's present form, may be run from any node on green.  For best
results have your ssh agent up and running. """

    parser = OptionParser(epilog=epilog, description=description, version=version, usage=usage)
    
    # Set the default resources to be used if none are specified.
    default_res = ( "nodes="+str(user_params["default_nodes"])+":ppn="+str(user_params["default_ppn"])+
                    ",walltime="+user_params["default_walltime"]+",pmem="+user_params["default_pmem"] )
    
    # Initialise some variables (good practice)
    use_mpiexec = True
    job_name = None
    
    parser.add_option("-r", "--resources", dest="res" , metavar="RESOURCE_LIST", 
                      default="default_res", type="string",
                      help="set the resources part of the job "
                           "script. e.g. nodes=1:ppn=2:s203,walltime=23:59:59")
    parser.add_option("-g", "--gen_script", dest="script_fname", default=None,
                      type="string",
                      help="generate the script using the provided file "
                           "name only - do not submit it to the queue.")
    parser.add_option("-n", "--job_name", default=None,
                     dest="job_name", type="string",
                     help="the name to be given to the job on the queue.")
    parser.add_option("--no_mpiexec", dest="use_mpiexec", default=True,
                  action="store_false",
                  help="don't prepend the default mpiexec to the job.")
    
    # If we have arguments the parse them.
    (opt, args) = parser.parse_args()

    if len(args)<1:
        parser.print_usage()
        sys.exit(1)    

    # Construct the string which will be our submission job.
    for i in range( len(args) -1 ):
        args[i] = args[i]+" "
    job = "".join(args)
    
    # Set the job name if required
    if (opt.job_name == None):
        opt.job_name = args[0]
        
    # Get our current working directory
    pwd = getcwd()
    
    # Do we just want to generate the script?
    # If so then open the requested output file for writing.
    # If not then generate a temporary file and write to that.
    if (opt.script_fname != None):
        script = open(opt.script_fname, 'w')
    else:
        script = NamedTemporaryFile(mode='w',dir=pwd,prefix='.tmp_run_mpi')
        
    # If any of the basic resources weren't specified then use predefined defaults.
    # Note that nodes MUST come first
    if "nodes" not in opt.res:
        opt.res = "nodes="+str(user_params["default_nodes"])+":ppn="+user_params['default_ppn']+","+res
        
    if "walltime" not in opt.res:
        opt.res += ",walltime="+user_params["default_walltime"]
    
    if "pmem" not in opt.res:
        opt.res += ",pmem="+user_params["default_pmem"]
        
    # Now write the submission script
    write_script(script, opt.job_name, pwd, opt.res, job, opt.use_mpiexec)
    
    # If requested - submit the job
    if opt.script_fname == None:
        status = sp.call("ssh green.ssi.swin.edu.au 'qsub "+script.name+" '", shell=True)
        if status !=0:
            print "Submission to the queue failed with exit status", status
            sys.exit(2)
    
    # Close the file.  If it is a temp file it will automatically be deleted.
    script.close()

if __name__ == '__main__':
    main(sys.argv[1:])
